#pragma once

#include "index/primary_key_index/primary_key_indexer.h"
#include "index/primary_key_index/primary_key_file_writer/sorted_primary_key_file_writer.h"
#include "index/deletion_index/deletion_index_indexer.h"
#include "util/hashmap/hash_map.h"
#include "util/status.h"
#include "util/pool/pool.h"

namespace hawking {
namespace indexlib {

class PrimaryKeyMemIndexer : public PrimaryKeyIndexer {
public:
    PrimaryKeyMemIndexer() = default;
    ~PrimaryKeyMemIndexer() = default;

    Status Init();
    Status AddDocument(uint64_t pk, DocId doc);
    Status Dump(
        const std::string& pk_dir,
        const std::unordered_map<DocId, DocId>& old2new_map);

    virtual Status Lookup(
        const google::protobuf::RepeatedField<int64_t>& pks,
        SegmentContext* segment_ctx) override;
    virtual ReturnValue<DocId> Lookup(uint64_t pk) override;

    size_t MemoryUse() const {
        return pool_->UsedBytes();
    }

    std::unordered_map<DocId, DocId> GenerateOldNewDocMap(
        std::shared_ptr<DeletionMapIndexer> deletion_indexer) const;

    virtual uint32_t DocCount() const override {
        return hash_map_->Size();
    }

    const HashMap<uint64_t, DocId>& Data() const {
        return *hash_map_;
    }

    class PrimaryKeyMemIterator : public PrimaryKeyIterator {
    public:
        PrimaryKeyMemIterator(const HashMap<uint64_t, DocId>& pk_doc_map) :
            pk_doc_map_iter_(pk_doc_map.CreateIterator()) {}
        virtual ~PrimaryKeyMemIterator() = default;

        virtual bool HasNext() const override {
            return pk_doc_map_iter_->HasNext();
        }

        virtual void Next() {}

        virtual void Value(uint64_t* pk, DocId* doc_id) const {
            current_value_ = &(pk_doc_map_iter_->Next());
            *pk = current_value_->first;
            *doc_id = current_value_->second;
        }

    private:
        HashMap<uint64_t, DocId>::Iterator pk_doc_map_iter_;
        std::pair<uint64_t, DocId>* current_value_ = nullptr;
    };

    virtual std::shared_ptr<PrimaryKeyIterator> CreateIterator() const override {
        return std::make_shared<PrimaryKeyMemIterator>(Data());
    }

private:
    std::shared_ptr<MMapAllocator> allocator_;
    std::unique_ptr<Pool> pool_;

    // pk <---> docid
    std::unique_ptr<HashMap<uint64_t, DocId>> hash_map_;
    std::shared_ptr<PrimaryKeyFileWriter<uint64_t>> pk_file_writer_;
};

}
}